import torch

batch_size = 1
seq_len = 3
input_size = 4
hidden_size = 4
num_layers = 1

idx2char = ["e","h","l","o"]
x_data = [1,0,2,2,3]
y_data = [3,1,2,3,2]

one_hot_lookup = [[1,0,0,0],
                  [0,1,0,0],
                  [0,0,1,0],
                  [0,0,0,1]]

x_one_hot = [one_hot_lookup[x] for x in x_data]

inputs = torch.Tensor(x_one_hot).view(-1,batch_size,input_size)
labels = torch.LongTensor(y_data).view(-1,1)


class Model(torch.nn.Module):
    def __init__(self,input_size,hidden_size,batch_size):
        super(Model,self).__init__()

        self.batch_size = batch_size
        self.inputsize = input_size
        self.hidden_size = hidden_size

        self.rnncell = torch.nn.RNNCell(input_size=self.inputsize,hidden_size=self.hidden_size)


    def forward(self,input,hidden):
        hidden = self.rnncell(input,hidden)
        return hidden


    def init_hidden(self):
        return torch.zeros(self.batch_size,self.hidden_size)


net = Model(input_size,hidden_size,batch_size)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(),lr = 0.1)
print(inputs,labels)

for epoch in range(15):
    loss = 0

    optimizer.zero_grad()
    hidden = net.init_hidden()
    print("predicted string:", end ="")

    for input,label in zip(inputs,labels):


        hidden = net(input,hidden)
        # print(hidden)
        loss += criterion(hidden,label)
        # _,idx =hidden.max(dim=1)
        idx = hidden.max(dim=1)[-1]
        # print(idx.item(),"==========")
        print(idx2char[idx.item()],end="")
    loss.backward()
    optimizer.step()

    print(",Epoch[%d/15] loss=%.4f" % (epoch+1, loss.item()))



